#!/usr/bin/env python3
"""
Author: Haoran Peng
Email: gavinsweden@gmail.com
"""

from heapq import heappop, heappush
from typing import Dict, List, Set, Tuple

import numpy as np
from scipy.spatial import KDTree

from .grid import Grid
from .neighbour_table import NeighbourTable
from .state import State


class Planner:
    def __init__(
        self, grid_size: int, robot_radius: int, static_obstacles: List[Tuple[int, int]]
    ):
        self.grid_size = grid_size
        self.robot_radius = robot_radius
        np_static_obstacles = np.array(static_obstacles)
        self.static_obstacles = KDTree(np_static_obstacles)

        # Make the grid according to the grid size
        self.grid = Grid(grid_size, np_static_obstacles)
        # Make a lookup table for looking up neighbours of a grid
        self.neighbour_table = NeighbourTable(self.grid.grid)

    """
    An admissible and consistent heuristic for A*
    """

    @staticmethod
    def h(start: np.ndarray, goal: np.ndarray) -> int:
        return int(np.linalg.norm(start - goal, 1))  # L1 norm

    @staticmethod
    def l2(start: np.ndarray, goal: np.ndarray) -> int:
        return int(np.linalg.norm(start - goal, 2))  # L2 norm

    """
    Check whether the nearest static obstacle is within radius
    """

    def safe_static(self, grid_pos: np.ndarray) -> bool:
        _, nn = self.static_obstacles.query(grid_pos)
        return self.l2(grid_pos, self.static_obstacles.data[nn]) > self.robot_radius

    """
    Space-Time A*
    """

    def plan(
        self,
        start: Tuple[int, int],
        goal: Tuple[int, int],
        dynamic_obstacles: Dict[int, Set[Tuple[int, int]]],
        speed_fix: Set[Tuple[int, int, int, int, int]],
        agent_id: int,
        semi_dynamic_obstacles: Dict[int, Set[Tuple[int, int]]] = None,
        max_iter: int = 500,
        debug: bool = False,
    ) -> np.ndarray:
        # Prepare dynamic obstacles
        dynamic_obstacles = dict(
            (k, np.array(list(v))) for k, v in dynamic_obstacles.items()
        )
        speed_space = [item for item in speed_fix]

        # Assume dynamic obstacles are agents with same radius, distance needs to be 2*radius
        def safe_dynamic(grid_pos: np.ndarray, time: int, bod: int) -> bool:
            nonlocal dynamic_obstacles
            return all(
                self.l2(grid_pos, obstacle) > bod * self.robot_radius
                for obstacle in dynamic_obstacles.setdefault(time, np.array([]))
            )

        def not_safe_speed(neighbour_pos: np.ndarray, current_pos: np.ndarray) -> bool:
            nonlocal speed_space
            # 这个函数需要排除速度区间内不满足条件的邻居,0为加速
            for space in speed_space:
                speed_type = space[0]
                start = space[1], space[2]
                end = space[3], space[4]

                if speed_type == 0 and start[0] <= current_pos[0] <= end[0]:
                    if not (
                        (
                            neighbour_pos[0] == current_pos[0] + 5 * self.grid_size
                            and neighbour_pos[1] == current_pos[1]
                        )
                        or (
                            neighbour_pos[0] == current_pos[0] + 2 * self.grid_size
                            and neighbour_pos[1] == current_pos[1] + 2 * self.grid_size
                        )
                        or (
                            neighbour_pos[0] == current_pos[0] + 2 * self.grid_size
                            and neighbour_pos[1] == current_pos[1] - 2 * self.grid_size
                        )
                    ):
                        return True
                elif speed_type == 3 and start[0] <= current_pos[0] <= end[0]:
                    if not (
                        (
                            neighbour_pos[0] == current_pos[0] + 1 * self.grid_size
                            and neighbour_pos[1] == current_pos[1]
                        )
                        or (
                            neighbour_pos[0] == current_pos[0]
                            and neighbour_pos[1] == current_pos[1] - 1 * self.grid_size
                        )
                        or (
                            neighbour_pos[0] == current_pos[0]
                            and neighbour_pos[1] == current_pos[1] + 1 * self.grid_size
                        )
                    ):
                        return True
                elif speed_type == 2 and start[0] <= current_pos[0] <= end[0]:
                    if not (
                        (
                            neighbour_pos[0] == current_pos[0] + 2 * self.grid_size
                            and neighbour_pos[1] == current_pos[1] + 2 * self.grid_size
                        )
                        or (
                            neighbour_pos[0] == current_pos[0] + 5 * self.grid_size
                            and neighbour_pos[1] == current_pos[1]
                        )
                    ):
                        return True
            return False

        # Prepare semi-dynamic obstacles, consider them static after specific timestamp
        if semi_dynamic_obstacles is None:
            semi_dynamic_obstacles = dict()
        else:
            semi_dynamic_obstacles = dict(
                (k, np.array(list(v))) for k, v in semi_dynamic_obstacles.items()
            )

        def safe_semi_dynamic(grid_pos: np.ndarray, time: int) -> bool:
            nonlocal semi_dynamic_obstacles
            for timestamp, obstacles in semi_dynamic_obstacles.items():
                flag = True
                if time >= timestamp:
                    flag = all(
                        self.l2(grid_pos, obstacle) > 2 * self.robot_radius
                        for obstacle in obstacles
                    )
                if not flag:
                    return False
            return True

        start = self.grid.snap_to_grid(np.array(start))
        goal = self.grid.snap_to_grid(np.array(goal))

        # Initialize the start state
        s = State(start, 0, 0, self.h(start, goal))

        open_set = [s]
        closed_set = set()

        # Keep track of parent nodes for reconstruction
        came_from = dict()

        iter_ = 0
        while open_set and iter_ < max_iter:
            iter_ += 1
            current_state = open_set[0]  # Smallest element in min-heap
            if current_state.pos_equal_to(goal):
                if debug:
                    print("STA*: Path found after {0} iterations".format(iter_))
                return self.reconstruct_path(came_from, current_state)

            closed_set.add(heappop(open_set))
            epoch = current_state.time + 1
            # if iter_ < agent_id*2:
            #     for neighbour in self.neighbour_table.lookup(current_state.pos):
            #         neighbour_state = State(neighbour, epoch, current_state.g_score + 1,self.h(neighbour, goal))
            #         if neighbour_state.pos[0] == current_state.pos[0] and neighbour_state.pos[1] == current_state.pos[1]-20:
            #             if neighbour_state not in open_set:
            #                 came_from[neighbour_state] = current_state
            #                 heappush(open_set, neighbour_state)
            #     continue

            for neighbour in self.neighbour_table.lookup(current_state.pos):
                if not_safe_speed(neighbour, current_state.pos):
                    continue
                # type_conflict = conflict_space
                # if type_conflict == "parallel" and location_type == "left":
                #     # 只留下(2,2),(2,-2),(5,0)

                #     if not (
                #         (
                #             neighbour[0] == current_state.pos[0] + 5 * self.grid_size
                #             and neighbour[1] == current_state.pos[1]
                #         )
                #         or (
                #             neighbour[0] == current_state.pos[0] + 2 * self.grid_size
                #             and neighbour[1]
                #             == current_state.pos[1] + 2 * self.grid_size
                #         )
                #         or (
                #             neighbour[0] == current_state.pos[0] + 2 * self.grid_size
                #             and neighbour[1]
                #             == current_state.pos[1] - 2 * self.grid_size
                #         )
                #     ):
                #         continue

                # elif type_conflict == "parallel" and location_type == "right":
                #     if not (
                #         (
                #             neighbour[0] == current_state.pos[0] + 1 * self.grid_size
                #             and neighbour[1] == current_state.pos[1]
                #         )
                #         or (
                #             neighbour[0] == current_state.pos[0]
                #             and neighbour[1]
                #             == current_state.pos[1] - 1 * self.grid_size
                #         )
                #         or (
                #             neighbour[0] == current_state.pos[0]
                #             and neighbour[1]
                #             == current_state.pos[1] + 1 * self.grid_size
                #         )
                #     ):
                #         continue
                # elif type_conflict == "point" and location_type == "left":
                #     if (
                #         (
                #             neighbour[0] == current_state.pos[0] + 5 * self.grid_size
                #             and neighbour[1] == current_state.pos[1]
                #         )
                #         or (
                #             neighbour[0] == current_state.pos[0] + 2 * self.grid_size
                #             and neighbour[1]
                #             == current_state.pos[1] + 2 * self.grid_size
                #         )
                #         or (
                #             neighbour[0] == current_state.pos[0] + 2 * self.grid_size
                #             and neighbour[1]
                #             == current_state.pos[1] - 2 * self.grid_size
                #         )
                #     ):
                #         continue
                # elif type_conflict == "point" and location_type == "right":
                #     if not (
                #         (
                #             neighbour[0] == current_state.pos[0] + 5 * self.grid_size
                #             and neighbour[1] == current_state.pos[1]
                #         )
                #         or (
                #             neighbour[0] == current_state.pos[0] + 2 * self.grid_size
                #             and neighbour[1]
                #             == current_state.pos[1] + 2 * self.grid_size
                #         )
                #         or (
                #             neighbour[0] == current_state.pos[0] + 2 * self.grid_size
                #             and neighbour[1]
                #             == current_state.pos[1] - 2 * self.grid_size
                #         )
                #     ):
                #         continue
                neighbour_state = State(
                    neighbour,
                    epoch,
                    current_state.g_score + 1,
                    self.h(neighbour, goal) + abs(neighbour[1] - goal[1]),
                )
                # Check if visited
                if neighbour_state in closed_set:
                    continue

                # Avoid obstacles
                if (
                    not self.safe_static(neighbour)
                    or not safe_dynamic(neighbour, epoch, 2)
                    or not safe_semi_dynamic(neighbour, epoch)
                ):
                    continue
                # # 提前避免冲突
                # if goal[1]<current_state.pos[1]:
                #     if neighbour_state.pos[0]!=current_state.pos[0]:
                #         continue

                # if not safe_dynamic(neighbour, epoch,5):
                #     if safe_dynamic_right_below(neighbour, epoch) and \
                #          goal[1]<neighbour_state.pos[1]:
                #             if neighbour_state.pos[0]!=current_state.pos[0]:
                #                 continue
                #             neighbour_state.f_score = 1

                # 避免垂直转向
                try:
                    if (
                        neighbour_state.pos[0] == current_state.pos[0]
                        and current_state.pos[1] == came_from[current_state].pos[1]
                        or neighbour_state.pos[1] == current_state.pos[1]
                        and current_state.pos[0] == came_from[current_state].pos[0]
                    ):
                        continue
                    if (
                        neighbour_state.pos[1] == came_from[current_state].pos[1]
                        and neighbour_state.pos[0] != came_from[current_state].pos[0]
                        and current_state.pos[1] != neighbour_state.pos[1]
                    ):
                        continue
                except:
                    pass

                # Add to open set
                if neighbour_state not in open_set:
                    came_from[neighbour_state] = current_state
                    heappush(open_set, neighbour_state)

        if debug:
            print("STA*: Open set is empty, no path found.")
        return np.array([])

    """
    Reconstruct path from A* search result
    """

    def reconstruct_path(
        self, came_from: Dict[State, State], current: State
    ) -> np.ndarray:
        total_path = [current.pos]
        while current in came_from.keys():
            current = came_from[current]
            total_path.append(current.pos)
        return np.array(total_path[::-1])
